package transform

import (
	"crypto/rand"
	mathRand "math/rand/v2"
	"net"
	"smart-flow/faker/proto"
	"smart-flow/faker/sys"
	"sync/atomic"
	"time"
)

type QuitFunc func()

type DataSender struct {
	fakerIp        string       // 虚拟IP
	v6             bool         // 是否IPv6
	rateDelta      int64        // 速率浮动系数
	udpPerUseCount int          // 每次使用多少个地址
	udpAddresses   []string     // Udp地址列表
	dataPool       []byte       // 只读数据池
	poolSize       int64        // 数据池大小
	blockCount     int64        // 每次发送的数据倍率(blockCount*poolSize)
	udpRateLimit   atomic.Int64 // 每秒限制速率(Bps)
	udpBaseRate    atomic.Int64 // 每秒基准速率(限制速率/perUseCount)
	udpRunRate     atomic.Int64 // 每秒运行速率
	paused         atomic.Bool  // 暂停
	stoped         atomic.Bool  // 停止
	fakeSocket     *sys.FakeSocket
	quitFunc       QuitFunc
}

func (ds *DataSender) Init(data *proto.FakerLoginResp, poolSize int64, fakerIp string, needV6 int, quitFunc QuitFunc) *DataSender {
	dataPool := make([]byte, poolSize) // 生成随机数据池
	if _, err := rand.Read(dataPool); err != nil {
		return nil
	}
	ds.fakerIp = fakerIp
	if needV6 == 1 {
		ds.v6 = true
	}
	ds.rateDelta = data.RateDelta
	ds.udpAddresses = data.UdpAddresses
	ds.dataPool = dataPool
	ds.poolSize = poolSize
	ds.blockCount = data.BlockCount
	ds.quitFunc = quitFunc
	ds.UpdatePerUseCount(data.PerUseCount)
	ds.UpdateRates(data.UdpRateLimit)
	ds.fakeSocket = new(sys.FakeSocket)
	return ds
}

func (ds *DataSender) UpdateRates(udpRateLimit int64) {
	length := len(ds.udpAddresses)
	if length > 0 && udpRateLimit > 0 && ds.udpPerUseCount > 0 {
		udpRateLimit = udpRateLimit * KBit / Byte2Bit
		ds.udpRateLimit.Store(udpRateLimit)
		ds.udpBaseRate.Store(udpRateLimit / int64(ds.udpPerUseCount))
	}
}

func (ds *DataSender) UpdatePerUseCount(perUseCount int) {
	ds.udpPerUseCount = max(perUseCount, min(perUseCount, len(ds.udpAddresses)))
}

func (ds *DataSender) UpdateRateDelta(reteDelta int64) {
	ds.rateDelta = reteDelta
}

func (ds *DataSender) Start() {
	length := len(ds.udpAddresses)
	if length > 0 && ds.udpRateLimit.Load() > 0 {
		go ds.startUdp(length)
	}
}

func (ds *DataSender) startUdp(length int) {
	i := 0
	availableRate, rate, baseRate := int64(0), int64(0), int64(0)
	rateLimit := ds.udpRateLimit.Load()
	for {
		baseRate = ds.udpBaseRate.Load()
		availableRate = rateLimit - ds.udpRunRate.Load()
		if availableRate < baseRate/2 { // 可用速率小于base/2不会运行, 应对随机总速率后的微增, 确保正常分片完成后再运行
			time.Sleep(time.Second)
			continue
		}
		rate = randomShared(baseRate, availableRate, ds.rateDelta)
		// fmt.Printf("udp 运行速率: %d Kbps, 速率限制: %d Kbps, 可用速率: %d Kbps\n", rate*Byte2Bit/KBit, rateLimit*Byte2Bit/KBit, availableRate*Byte2Bit/KBit)
		ds.udpRunRate.Add(rate)
		go ds.sendData2Udp(ds.udpAddresses[i], rate)
		i++
		if i%ds.udpPerUseCount == 0 {
			rateLimit = ds.udpRateLimit.Load()
			rateLimit = randomRate(rateLimit, ds.rateDelta)
		}
		if i == length {
			i = 0
		}
	}
}

func (ds *DataSender) sendData2Udp(address string, rate int64) {
	defer ds.udpRunRate.Add(-rate)
	ip, port, err := net.SplitHostPort(address)
	if err != nil {
		ds.quitFunc()
		return
	}
	if !ds.fakeSocket.Open(ds.v6, ds.fakerIp, ip, port) {
		ds.quitFunc()
		return
	}
	defer ds.fakeSocket.Close()
	totalSize := ds.poolSize * ds.blockCount / int64(ds.udpPerUseCount) // 保证每次在不长的时间内完成(30s~3m30s)
	// fmt.Printf("udp 任务内共需发送: %d KB, 端口:%d, 速率限制: %d Kbps\n", totalSize/KB, port, rate*Byte2Bit/KBit)
	totalSent := int64(0)
	bytesInSec := int64(0)
	lastTime := time.Now()
	for totalSent < totalSize {
		if ds.Stoped() {
			return
		}
		if ds.Paused() {
			time.Sleep(time.Second)
			continue
		}
		now := time.Now()
		elapsed := now.Sub(lastTime)
		if elapsed >= time.Second {
			lastTime = now
			// fmt.Printf("udp 已发送 %d KB，端口:%d, 速率: %d Kbps, 速率限制: %d Kbps\n", totalSent/KB, port, bytesInSec*Byte2Bit/KBit, rate*Byte2Bit/KBit)
			bytesInSec = 0 // 每秒更新一次计数器
		}
		toWrite := rate - bytesInSec // 计算本次发送的数据量
		if toWrite <= 0 {
			time.Sleep(time.Second - elapsed) // 如果当前秒已经用完配额，则等待下一轮
			continue
		}
		packetSize := min(500, toWrite) // 确定单次发送的数据大小
		startPos := totalSent % ds.poolSize
		endPos := startPos + packetSize
		if endPos > ds.poolSize {
			startPos = 0
			endPos = packetSize
		}
		_, err = ds.fakeSocket.Send(ds.dataPool[startPos:endPos]) // 发送UDP数据
		if err != nil {
			ds.quitFunc()
			return
		}
		totalSent += packetSize
		bytesInSec += packetSize
	}
}

func (ds *DataSender) Pause() {
	ds.paused.Store(true)
}

func (ds *DataSender) Resume() {
	ds.paused.Store(false)
}

func (ds *DataSender) Paused() bool {
	return ds.paused.Load()
}

func (ds *DataSender) Stop() {
	ds.stoped.Store(true)
}

func (ds *DataSender) Stoped() bool {
	return ds.stoped.Load()
}

func randomShared(base, availableRate, rateDelta int64) int64 {
	var rate int64
	maxRate := base * (100 + rateDelta) / 100
	if availableRate <= maxRate+base/2 { // 确保不会出现很小的速率, 因为小于base/2不会执行
		rate = availableRate
	} else {
		deltaPercent := mathRand.Int64N(rateDelta + 1)
		rate = min(availableRate, base*(100+deltaPercent)/100) // base~base(100+rateDelta)%
	}
	return rate
}

func randomRate(rateLimit, rateDelta int64) int64 {
	if rateDelta <= 0 {
		return rateLimit
	}
	deltaPercent := mathRand.Int64N(rateDelta*2+1) - rateDelta
	return rateLimit * (100 + deltaPercent) / 100
}
